import torch
from sklearn.metrics import roc_auc_score
import logging
from tqdm import tqdm


def train_model(model, train_loader, val_loader, criterion, optimizer, device, epochs=5):
    """模型训练函数"""
    model.to(device)
    best_auc = 0.0

    for epoch in range(epochs):
        # 训练阶段
        model.train()
        train_loss = 0.0
        train_pbar = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{epochs} - Training")

        for cat_feats, num_feats, targets in train_pbar:
            cat_feats = cat_feats.to(device)
            num_feats = num_feats.to(device)
            targets = targets.to(device).view(-1, 1)

            optimizer.zero_grad()

            outputs = model(cat_feats, num_feats)
            loss = criterion(outputs, targets)

            # 损失函数反馈
            loss.backward()
            optimizer.step()

            train_loss += loss.item() * cat_feats.size(0)
            train_pbar.set_postfix({"batch_loss": loss.item()})

        # 计算训练集平均损失
        train_loss_avg = train_loss / len(train_loader.dataset)

        # 验证阶段
        model.eval()
        val_loss = 0.0
        all_preds, all_targets = [], []
        val_pbar = tqdm(val_loader, desc=f"Epoch {epoch + 1}/{epochs} - Validation")

        with torch.no_grad():
            for cat_feats, num_feats, targets in val_pbar:
                cat_feats = cat_feats.to(device)
                num_feats = num_feats.to(device)
                targets = targets.to(device).view(-1, 1)

                outputs = model(cat_feats, num_feats)
                loss = criterion(outputs, targets)

                val_loss += loss.item() * cat_feats.size(0)
                all_preds.extend(outputs.cpu().numpy())
                all_targets.extend(targets.cpu().numpy())

        # 计算验证集指标
        val_loss_avg = val_loss / len(val_loader.dataset)
        val_auc = roc_auc_score(all_targets, all_preds)

        # 日志输出
        logging.info(f"\nEpoch {epoch + 1} Summary:")
        logging.info(f"Train Loss: {train_loss_avg:.4f} | Val Loss: {val_loss_avg:.4f} | Val AUC: {val_auc:.4f}")

        # 保存最优模型
        if val_auc > best_auc:
            best_auc = val_auc
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'best_auc': best_auc
            }, '../res/deepfm_best.pth')
            logging.info(f"Saved best model with AUC: {best_auc:.4f}")

    return model
